{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X_train.shape= (671, 60, 1)\n",
      "y_train.shape= (671, 1)\n"
     ]
    }
   ],
   "source": [
    "# 从CSV文件中导入数据\n",
    "data = pd.read_csv(\"./zgpa_train.csv\")\n",
    "\n",
    "dataset_train = data.sort_values(by='date').reset_index(drop=True)\n",
    "training_set = dataset_train.iloc[:, 1:2].values\n",
    "# 数据归一化\n",
    "max = np.max(training_set)\n",
    "min = np.min(training_set)\n",
    "training_set_norm = []\n",
    "for i in range(training_set.shape[0]):\n",
    "    training_set[i] = (training_set[i]-min) / (max-min)\n",
    "training_set_norm = training_set\n",
    "# 每条样本含60个时间步，对应下一时间步的标签值\n",
    "X_train = []\n",
    "y_train = []\n",
    "time_steps = 60\n",
    "for i in range(time_steps, training_set.shape[0]):\n",
    "    X_train.append(training_set_norm[i-time_steps:i, 0])\n",
    "    y_train.append(training_set_norm[i, 0])\n",
    "X_train, y_train = np.array(X_train), np.array(y_train).reshape(-1, 1)\n",
    "X_train = X_train.reshape(X_train.shape[0], -1, 1)\n",
    "print(f'{X_train.shape= }')\n",
    "print(f'{y_train.shape= }')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 实现一个LSTM网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM:\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        np.random.seed(0)\n",
    "\n",
    "        # forget gate weight and bias\n",
    "        self.wxf = np.random.randn(input_size, hidden_size)\n",
    "        self.whf = np.random.randn(hidden_size, hidden_size)\n",
    "        self.bf = np.zeros((1, hidden_size))\n",
    "\n",
    "        # update gate weight and bias\n",
    "        self.wxg = np.random.randn(input_size, hidden_size)\n",
    "        self.whg = np.random.randn(hidden_size, hidden_size)\n",
    "        self.bg = np.zeros((1, hidden_size))\n",
    "\n",
    "        self.wxi = np.random.randn(input_size, hidden_size)\n",
    "        self.whi = np.random.randn(hidden_size, hidden_size)\n",
    "        self.bi = np.zeros((1, hidden_size))\n",
    "\n",
    "        # output gate weight and bias\n",
    "        self.wxo = np.random.randn(input_size, hidden_size)\n",
    "        self.who = np.random.randn(hidden_size, hidden_size)\n",
    "        self.bo = np.zeros((1, hidden_size))\n",
    "\n",
    "        # last output weight and bias\n",
    "        self.why = np.random.randn(hidden_size, output_size)\n",
    "        self.by = np.zeros((1, output_size))\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.x = x\n",
    "        batch_size, time_steps, _ = x.shape\n",
    "\n",
    "        h = np.zeros((batch_size, self.hidden_size))  # 初始化隐藏层状态\n",
    "        c = np.zeros((batch_size, self.hidden_size))  # 初始化记忆细胞\n",
    "        self.h_temp, self.c_temp, self.f_temp, self.i_temp, self.g_temp, self.o_temp = [\n",
    "        ], [], [], [], [], []\n",
    "        self.h_temp.append(h)\n",
    "        self.c_temp.append(c)\n",
    "\n",
    "        self.t = time_steps  # times_steps\n",
    "        for i in range(self.t):\n",
    "\n",
    "            x_t = x[:, i, :]\n",
    "            f = self.sigmoid(x_t @ self.wxf + h @ self.whf+self.bf)\n",
    "            self.f_temp.append(f)\n",
    "            g = self.tanh(x_t @ self.wxg + h @ self.whg + self.bg)\n",
    "            self.g_temp.append(g)\n",
    "            i = self.sigmoid(x_t @ self.wxi+h @ self.whi + self.bi)\n",
    "            self.i_temp.append(i)\n",
    "            c = c * f + g * i\n",
    "            self.c_temp.append(c)\n",
    "            o = self.sigmoid(x_t @ self.wxo + h @ self.who + self.bo)\n",
    "            self.o_temp.append(o)\n",
    "            m = self.tanh(c)\n",
    "\n",
    "            h = o * m\n",
    "            self.h_temp.append(h)\n",
    "        y = h @ self.why + self.by\n",
    "\n",
    "        return y\n",
    "\n",
    "    def backward(self, y_true, y_pred, lr):\n",
    "        # dL/dy\n",
    "        dy = y_pred-y_true\n",
    "\n",
    "        # dL/dby\n",
    "        dby = np.sum(dy, axis=0, keepdims=True)\n",
    "\n",
    "        # dL/dwhy\n",
    "        dwhy = self.h_temp[-1].T @ dy\n",
    "\n",
    "        # dL/dh\n",
    "        dh = dy @ self.why.T\n",
    "\n",
    "        dwxf = np.zeros_like(self.wxf)\n",
    "        dwhf = np.zeros_like(self.whf)\n",
    "        dbf = np.zeros_like(self.bf)\n",
    "\n",
    "        dwxg = np.zeros_like(self.wxg)\n",
    "        dwhg = np.zeros_like(self.whg)\n",
    "        dbg = np.zeros_like(self.bg)\n",
    "\n",
    "        dwxi = np.zeros_like(self.wxi)\n",
    "        dwhi = np.zeros_like(self.whi)\n",
    "        dbi = np.zeros_like(self.bi)\n",
    "\n",
    "        dwxo = np.zeros_like(self.wxo)\n",
    "        dwho = np.zeros_like(self.who)\n",
    "        dbo = np.zeros_like(self.bo)\n",
    "        for i in range(self.t)[::-1]:\n",
    "            temp_o = dh * \\\n",
    "                self.tanh(self.c_temp[i+1]) * \\\n",
    "                self.o_temp[i] * (1-self.o_temp[i])\n",
    "\n",
    "            # dL/dwxo\n",
    "            dwxo += self.x[:, i, :].T @ temp_o\n",
    "\n",
    "            # dL/dwho\n",
    "            dwho += self.h_temp[i].T @ temp_o\n",
    "\n",
    "            # dL/dbo\n",
    "            dbo += np.sum(temp_o, axis=0, keepdims=True)\n",
    "\n",
    "            # dL/dc\n",
    "            dc = dh * self.o_temp[i] * (1-self.c_temp[i+1]**2)\n",
    "\n",
    "            temp_f = dc * self.c_temp[i] * self.f_temp[i] * (1-self.f_temp[i])\n",
    "\n",
    "            # dL/dwxf\n",
    "            dwxf += self.x[:, i, :].T @ temp_f\n",
    "\n",
    "            # dL/dwhf\n",
    "            dwhf += self.h_temp[i].T @ temp_f\n",
    "\n",
    "            # dL/dbf\n",
    "            dbf += np.sum(temp_f, axis=0, keepdims=True)\n",
    "\n",
    "            temp_i = dc * self.g_temp[i] * self.i_temp[i] * (1-self.i_temp[i])\n",
    "\n",
    "            # dL/dwxi\n",
    "            dwxi += self.x[:, i, :].T @ temp_i\n",
    "\n",
    "            # dL/dwhi\n",
    "            dwhi += self.h_temp[i].T @ temp_i\n",
    "\n",
    "            # dL/dbi\n",
    "            dbi += np.sum(temp_i, axis=0, keepdims=True)\n",
    "\n",
    "            temp_g = dc * self.i_temp[i] * (1-self.g_temp[i]**2)\n",
    "\n",
    "            # dL/dwxg\n",
    "            dwxg += self.x[:, i, :].T @ temp_g\n",
    "\n",
    "            # dL/dwhg\n",
    "            dwhg += self.h_temp[i].T @ temp_g\n",
    "\n",
    "            # dL/dbg\n",
    "            dbg += np.sum(temp_g, axis=0, keepdims=True)\n",
    "\n",
    "            # dL/dh\n",
    "            dh = temp_o @ self.who.T + temp_f @ self.whf.T + \\\n",
    "                temp_i @ self.whi.T + temp_g @ self.whg.T\n",
    "\n",
    "            # dL/dc\n",
    "\n",
    "            dc = dc * self.o_temp[i-1] * self.tanh_derivative(self.c_temp[i]) * \\\n",
    "            (self.f_temp[i] + \\\n",
    "            self.c_temp[i] * self.f_temp[i] * (1-self.f_temp[i]) @ self.whf.T  + \\\n",
    "            self.i_temp[i] * (1-self.g_temp[i]**2) @ self.whg.T + \\\n",
    "            self.g_temp[i] * self.i_temp[i] * (1-self.i_temp[i]) @ self.whi.T\n",
    "            )\n",
    "        \n",
    "        #update forget gate weight and bias\n",
    "        self.wxf -= lr * dwxf\n",
    "        self.whf -= lr * dwhf\n",
    "        self.bf -= lr * dbf\n",
    "\n",
    "        # update update gate weight and bias\n",
    "        self.wxg -= lr * dwxg\n",
    "        self.whg -= lr * dwhg\n",
    "        self.bg -= lr * dbg\n",
    "\n",
    "        self.wxi -= lr * dwxi\n",
    "        self.whi -= lr * dwhi\n",
    "        self.bi -= lr * dbi\n",
    "\n",
    "        # update output gate weight and bias\n",
    "        self.wxo -= lr * dwxo\n",
    "        self.who -= lr * dwho\n",
    "        self.bo -= lr * dbo\n",
    "\n",
    "        # update last output weight and bias\n",
    "        self.why -= lr * dwhy\n",
    "        self.by -= lr * dby\n",
    "\n",
    "    def sigmoid(self, x):\n",
    "        return 1 / (1 + np.exp(-x))\n",
    "\n",
    "    def sigmoid_derivative(self, x):\n",
    "        return self.sigmoid(x) * (1 - self.sigmoid(x))\n",
    "\n",
    "    def tanh(self, x):\n",
    "        return np.tanh(x)\n",
    "\n",
    "    def tanh_derivative(self, x):\n",
    "        return 1 - self.tanh(x)**2\n",
    "\n",
    "    def mean_squared_error_loss(self, y_pred, y_true):\n",
    "        \"\"\"\n",
    "        均方误差损失函数\n",
    "        :param y_pred : 预测值 \n",
    "        :param y_true : 真实值 \n",
    "        :returns loss : 损失值 \n",
    "        \"\"\"\n",
    "        loss = np.sum((y_pred-y_true)**2) / len(y_true)\n",
    "        return loss\n",
    "\n",
    "    def train(self, x, y, epochs, batch_size, lr):\n",
    "        self.losses = []\n",
    "        for epoch in range(epochs):\n",
    "            mini_x_trains = [x[k:k+batch_size] for k in range(0, x.shape[0], batch_size)]\n",
    "            mini_y_labels = [y[k:k+batch_size] for k in range(0, x.shape[0], batch_size)]\n",
    "            for mini_x_train, mini_y_label in zip(mini_x_trains, mini_y_labels):\n",
    "                output = self.forward(mini_x_train)\n",
    "\n",
    "                self.backward(mini_y_label, output, lr)\n",
    "\n",
    "            output = self.forward(x)\n",
    "            loss = self.mean_squared_error_loss(output, y)\n",
    "            # 记录所有损失值\n",
    "            self.losses.append(loss)\n",
    "            if (epoch+1) % 10 == 0:\n",
    "                print(f'第{epoch+1}轮, loss = {loss}')\n",
    "    def predict(self,x):\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "第10轮, loss = 0.0005105502007662448\n",
      "第20轮, loss = 0.00036682844687137533\n",
      "第30轮, loss = 0.00035315113548982813\n",
      "第40轮, loss = 0.0003469640337890438\n",
      "第50轮, loss = 0.0003426206271589265\n",
      "第60轮, loss = 0.0003389899292494923\n",
      "第70轮, loss = 0.0003356993850608485\n",
      "第80轮, loss = 0.00033260627258277416\n",
      "第90轮, loss = 0.0003296505398418044\n",
      "第100轮, loss = 0.0003268044338569196\n"
     ]
    }
   ],
   "source": [
    "model = LSTM(1, 2, 1)\n",
    "model.train(X_train[0:400],y_train[0:400],100,100,0.01)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
